import soundfile as sf
import torch
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 加载模型并指定设备
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-Omni-7B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    load_in_4bit=True
)

model.disable_talker()

processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")

conversation = [
    {
        "role": "system",
        "content": "你是一个视频异常行为预测以及原因解释模型，请你回答每一条输入的异常行为预测和原因解释，以Python元组的形式返回。（例如：(摔倒,没有站稳)）",
    },
    {
        "role": "user",
        "content": [
            {"type": "video", "video": r"D:\Python\Projects\TeamProjects\health-care\pymodel\data\fall.mp4"},
        ],
    },
]

# 设置视频中是否使用音频
USE_AUDIO_IN_VIDEO = True

text = processor.apply_chat_template(
    conversation, add_generation_prompt=True, tokenize=False)

audios, images, videos = process_mm_info(
    conversation, use_audio_in_video=USE_AUDIO_IN_VIDEO)

# 处理输入并移动到正确设备
inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt",
                   padding=True, use_audio_in_video=USE_AUDIO_IN_VIDEO)

# 将所有输入移动到与模型相同的设备
inputs = {k: v.to(device) if isinstance(v, torch.Tensor)
          else v for k, v in inputs.items()}

# 生成文本
text_ids = model.generate(**inputs, return_audio=False)

text = processor.batch_decode(
    text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)

print(text)
